import torch
import torch.nn as nn
from pytorch_pretrained_bert import BertModel


class BERT(nn.Module):
    def __init__(self, config):
        super(BERT, self).__init__()

        self.bert = BertModel.from_pretrained(config.bert_path)

        self.output = nn.Linear(768, 23)
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        out = self.bert(x, output_all_encoded_layers=False)[1]

        out = self.output(out)
        out = self.dropout(out)

        return out

    def predict(self, x):
        return torch.argmax(self.forward(x), dim=1)